from langgraph.graph import END, StateGraph
from typing import Dict, TypedDict, List, Any, Literal, Optional
import re
import json
import random

# 定义工具结构
class Tool(TypedDict):
    name: str
    description: str
    parameters: Dict[str, Any]

# 定义大模型响应结构
class ModelResponse(TypedDict):
    thought: str
    tool: Optional[str]  # 可能为None，表示不需要工具
    parameters: Dict[str, Any]

# 定义状态结构
class AgentState(TypedDict):
    user_input: str           # 用户输入的字符串
    model_response: ModelResponse  # 大模型结构化响应
    tool_result: Any          # 工具执行结果
    final_output: str         # 最终输出给用户的内容
    conversation: List[Dict]  # 对话历史

# 模拟可用工具
TOOLS: List[Tool] = [
    {
        "name": "add",
        "description": "执行两个数字的加法运算",
        "parameters": {
            "type": "object",
            "properties": {
                "a": {"type": "number", "description": "第一个数字"},
                "b": {"type": "number", "description": "第二个数字"}
            },
            "required": ["a", "b"]
        }
    },
    {
        "name": "subtract",
        "description": "执行两个数字的减法运算",
        "parameters": {
            "type": "object",
            "properties": {
                "a": {"type": "number", "description": "被减数"},
                "b": {"type": "number", "description": "减数"}
            },
            "required": ["a", "b"]
        }
    }
]

# 模拟大模型对话节点
def model_dialogue_node(state: AgentState) -> Dict:
    """模拟大模型对话节点，分析用户输入并决定下一步行动"""
    print(f"\n💬 对话历史: {state.get('conversation', [])[-1:] if state.get('conversation') else '无'}")
    print(f"🧠 用户输入: '{state['user_input']}'")
    print(f"🛠️ 可用工具: {[t['name'] for t in TOOLS]}")

    # 模拟大模型思考过程
    thought = ""
    tool_name = None
    parameters = {}

    # 提取数字
    numbers = re.findall(r'-?\d+', state["user_input"])

    # 根据输入内容决定响应
    if "加" in state["user_input"] or "+" in state["user_input"] or "和" in state["user_input"]:
        thought = "用户要求进行加法运算"
        tool_name = "add"
    elif "减" in state["user_input"] or "-" in state["user_input"] or "差" in state["user_input"]:
        thought = "用户要求进行减法运算"
        tool_name = "subtract"
    elif any(char.isdigit() for char in state["user_input"]):
        # 包含数字但没有明确操作 - 需要澄清
        thought = "用户提供了数字但未明确操作类型，需要澄清"
        tool_name = None
    else:
        # 非数学问题
        thought = "这不是数学问题，将直接回复用户"
        tool_name = None

    # 处理需要工具的情况
    if tool_name:
        # 确保提取到两个数字
        if len(numbers) >= 2:
            try:
                a = int(numbers[0])
                b = int(numbers[1])
                parameters = {"a": a, "b": b}
                thought += f"。将使用数字 {a} 和 {b} 进行计算"
            except ValueError:
                thought += "。数字解析失败，将请求用户提供有效数字"
                tool_name = None
        else:
            thought += "。输入中缺少足够数字，将请求用户提供数字"
            tool_name = None

    # 构建结构化响应
    response: ModelResponse = {
        "thought": thought,
        "tool": tool_name,
        "parameters": parameters
    }

    # 添加随机性使模拟更真实
    if random.random() < 0.2:  # 20%概率添加额外思考
        response["thought"] += "（额外思考：用户可能需要更多帮助）"

    print(f"🤖 模拟大模型响应: {json.dumps(response, indent=2, ensure_ascii=False)}")

    # 更新对话历史
    conversation = state.get("conversation", [])
    conversation.append({
        "role": "assistant",
        "content": thought,
        "tool": tool_name
    })

    return {
        "model_response": response,
        "conversation": conversation
    }

# 工具执行函数
def execute_tool(state: AgentState) -> Dict:
    """根据模型选择的工具执行相应操作"""
    if not state["model_response"]["tool"]:
        return {}  # 没有工具需要执行

    tool_name = state["model_response"]["tool"]
    params = state["model_response"]["parameters"]
    a = params.get("a", 0)
    b = params.get("b", 0)

    print(f"🔧 执行工具: {tool_name}(a={a}, b={b})")

    if tool_name == "add":
        result = a + b
        print(f"➕ 加法结果: {a} + {b} = {result}")
    elif tool_name == "subtract":
        result = a - b
        print(f"➖ 减法结果: {a} - {b} = {result}")
    else:
        result = None
        print(f"⚠️ 未知工具: {tool_name}")

    return {"tool_result": result}

# 生成最终响应
def generate_final_response(state: AgentState) -> Dict:
    """生成最终输出给用户的内容"""
    tool_name = state["model_response"]["tool"]
    params = state["model_response"]["parameters"]
    result = state.get("tool_result")
    thought = state["model_response"]["thought"]

    # 根据情况生成不同响应
    if not tool_name:
        # 不需要工具的情况
        if "需要澄清" in thought:
            output = "您想对这些数字进行加法还是减法运算？"
        elif "直接回复" in thought:
            output = "您好！我可以帮助您进行加法和减法运算。请告诉我您想计算什么？"
        else:
            output = "我无法理解您的请求，请提供更清晰的数学问题。"
    elif result is None:
        output = "抱歉，计算过程中出现了问题。"
    elif tool_name == "add":
        output = f"计算结果: {params['a']} + {params['b']} = {result}"
    elif tool_name == "subtract":
        output = f"计算结果: {params['a']} - {params['b']} = {result}"
    else:
        output = f"计算结果: {result}"

    # 添加思考过程（简化版）
    output = f"{output}"

    # 更新对话历史
    conversation = state.get("conversation", [])
    conversation.append({
        "role": "assistant",
        "content": output
    })

    print(f"📝 最终输出: {output}")
    return {
        "final_output": output,
        "conversation": conversation
    }

# 添加条件边 - 根据模型决策路由
def should_use_tool(state: AgentState) -> Literal["use_tool", "no_tool"]:
    """根据模型响应决定是否使用工具"""
    if state["model_response"]["tool"]:
        print("⤵️ 路由决策: 需要执行工具")
        return "use_tool"
    else:
        print("⤵️ 路由决策: 直接生成响应")
        return "no_tool"

# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("model_dialogue", model_dialogue_node)
    workflow.add_node("execute_tool", execute_tool)
    workflow.add_node("generate_response", generate_final_response)

    # 设置入口点
    workflow.set_entry_point("model_dialogue")

    # 条件路由
    workflow.add_conditional_edges(
        "model_dialogue",
        should_use_tool,
        {
            "use_tool": "execute_tool",
            "no_tool": "generate_response"
        }
    )

    # 连接工具执行后的节点
    workflow.add_edge("execute_tool", "generate_response")

    # 连接到结束
    workflow.add_edge("generate_response", END)

    return workflow.compile()

# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 测试案例
    test_cases = [
        {"input": "5加3等于多少？", "expected": ["5 + 3 = 8"]},
        {"input": "10减去7", "expected": ["10 - 7 = 3"]},
        {"input": "计算20和15的和", "expected": ["20 + 15 = 35"]},
        {"input": "100 - 50", "expected": ["100 - 50 = 50"]},
        {"input": "8和4的差是多少？", "expected": ["8 - 4 = 4"]},
        {"input": "3和5", "expected": ["加法", "减法", "澄清"]},  # 需要澄清
        {"input": "你好", "expected": ["帮助", "加法", "减法"]},  # 问候语
        {"input": "42和100", "expected": ["加法", "减法", "澄清"]},  # 需要澄清
        {"input": "计算圆的面积", "expected": ["无法理解", "帮助"]},  # 非数学问题
    ]

    for test in test_cases:
        print("\n" + "="*50)
        print(f"测试输入: '{test['input']}'")

        # 初始状态
        state = {
            "user_input": test["input"],
            "conversation": [
                {"role": "user", "content": test["input"]}
            ]
        }

        result = app.invoke(state)

        # 验证结果
        final_output = result["final_output"]
        passed = any(exp in final_output for exp in test["expected"])

        if passed:
            print(f"✅ 测试通过! 输出: '{final_output}'")
        else:
            print(f"❌ 测试失败! 期望包含: {test['expected']}，实际: '{final_output}'")

        # 打印对话历史
        print("\n🗣️ 完整对话历史:")
        for msg in result["conversation"]:
            role = msg["role"].capitalize()
            content = msg["content"]
            tool = msg.get("tool", "")
            print(f"{role}: {content} {f'(工具: {tool})' if tool else ''}")
